/*
 * Copyright (C) 2010 TopCoder Inc., All Rights Reserved.
 */
package com.hercules.xre.protocol.testing.netty.websocket;

import java.security.MessageDigest;

import org.jboss.netty.buffer.ChannelBuffer;
import org.jboss.netty.buffer.ChannelBuffers;
import org.jboss.netty.channel.ChannelFuture;
import org.jboss.netty.channel.ChannelFutureListener;
import org.jboss.netty.channel.ChannelHandlerContext;
import org.jboss.netty.channel.ChannelPipeline;
import org.jboss.netty.channel.Channels;
import org.jboss.netty.channel.ExceptionEvent;
import org.jboss.netty.channel.MessageEvent;
import org.jboss.netty.channel.SimpleChannelUpstreamHandler;
import org.jboss.netty.handler.codec.http.DefaultHttpResponse;
import org.jboss.netty.handler.codec.http.HttpHeaders;
import org.jboss.netty.handler.codec.http.HttpMethod;
import org.jboss.netty.handler.codec.http.HttpRequest;
import org.jboss.netty.handler.codec.http.HttpResponse;
import org.jboss.netty.handler.codec.http.HttpResponseStatus;
import org.jboss.netty.handler.codec.http.HttpVersion;
import org.jboss.netty.handler.codec.http.websocket.WebSocketFrame;
import org.jboss.netty.handler.codec.http.websocket.WebSocketFrameDecoder;
import org.jboss.netty.handler.codec.http.websocket.WebSocketFrameEncoder;
import org.jboss.netty.util.CharsetUtil;

/**
 * This class mainly handles the handshaking of the web socket communication.
 * 
 * @author TCSASSEMBLER
 * @version 1.0
 */
public class WebSocketServerHandler extends SimpleChannelUpstreamHandler {
    /**
     * This path of the web socket.
     */
    private static final String WEBSOCKET_PATH = "/xre";

    /**
     * <p>
     * Method called when message has been received and processed. It will
     * response the HTTP request (handshaking). WebSocketFrame messages are
     * simply sent to the upstream.
     * </p>
     * 
     * @param e
     *            message event, can't be null
     * @param ctx
     *            channel handler context, can't be null
     * @throws Exception
     *             if any error occurs
     */
    @Override
    public void messageReceived(ChannelHandlerContext ctx, MessageEvent e) throws Exception {
        Object msg = e.getMessage();
        if (msg instanceof HttpRequest) {
            handleHttpRequest(ctx, (HttpRequest) msg);
        } else if (msg instanceof WebSocketFrame) {
            // directly fire this message to the up stream
            Channels.fireMessageReceived(ctx, msg);
        }
    }

    /**
     * <p>
     * Responses the HTTP request (handshaking). It will replace the HTTP
     * decoder/encoder with WebSocketFrame decoder/encoder after this is done.
     * </p>
     * 
     * @param req
     *            the HttpRequest
     * @param ctx
     *            channel handler context, can't be null
     */
    private void handleHttpRequest(ChannelHandlerContext ctx, HttpRequest req) throws Exception {
        // Allow only GET methods.
        if (req.getMethod() != HttpMethod.GET) {
            sendHttpResponse(ctx, req, new DefaultHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.FORBIDDEN));
            return;
        }

        // Serve the WebSocket connection.
        if (req.getUri().equals(WEBSOCKET_PATH)) {
            if (HttpHeaders.Values.UPGRADE.equalsIgnoreCase(req.getHeader(HttpHeaders.Names.CONNECTION))
                    && HttpHeaders.Values.WEBSOCKET.equalsIgnoreCase(req.getHeader(HttpHeaders.Names.UPGRADE))) {
                HttpResponse res = new DefaultHttpResponse(HttpVersion.HTTP_1_1, new HttpResponseStatus(101,
                        "Web Socket Protocol Handshake"));
                res.addHeader(HttpHeaders.Names.UPGRADE, HttpHeaders.Values.WEBSOCKET);
                res.addHeader(HttpHeaders.Names.CONNECTION, HttpHeaders.Values.UPGRADE);
                
                // Fill in the headers and contents depending on handshake method.
                if (req.containsHeader(HttpHeaders.Names.SEC_WEBSOCKET_KEY1) &&
                    req.containsHeader(HttpHeaders.Names.SEC_WEBSOCKET_KEY2)) {
                    // New handshake method with a challenge
                    res.addHeader(HttpHeaders.Names.SEC_WEBSOCKET_ORIGIN, req.getHeader(HttpHeaders.Names.ORIGIN));
                    res.addHeader(HttpHeaders.Names.SEC_WEBSOCKET_LOCATION, getWebSocketLocation(req));
                    String protocol = req.getHeader(HttpHeaders.Names.SEC_WEBSOCKET_PROTOCOL);
                    if (protocol != null) {
                        res.addHeader(HttpHeaders.Names.SEC_WEBSOCKET_PROTOCOL, protocol);
                    }
                    
                    // Calculate the answer of the challenge.
                    String key1 = req.getHeader(HttpHeaders.Names.SEC_WEBSOCKET_KEY1);
                    String key2 = req.getHeader(HttpHeaders.Names.SEC_WEBSOCKET_KEY2);
                    int a = (int) (Long.parseLong(key1.replaceAll("[^0-9]", ""))
                            / key1.replaceAll("[^ ]", "").length());
                    int b = (int) (Long.parseLong(key2.replaceAll("[^0-9]", ""))
                            / key2.replaceAll("[^ ]", "").length());
                    long c = req.getContent().readLong();
                    ChannelBuffer input = ChannelBuffers.buffer(16);
                    input.writeInt(a);
                    input.writeInt(b);
                    input.writeLong(c);
                    ChannelBuffer output = ChannelBuffers.wrappedBuffer(MessageDigest.getInstance("MD5").digest(
                            input.array()));
                    res.setContent(output);
                } else {
                    // Old handshake method with no challenge
                    res.addHeader(HttpHeaders.Names.WEBSOCKET_ORIGIN, req.getHeader(HttpHeaders.Names.ORIGIN));
                    res.addHeader(HttpHeaders.Names.WEBSOCKET_LOCATION, getWebSocketLocation(req));
                    String protocol = req.getHeader(HttpHeaders.Names.WEBSOCKET_PROTOCOL);
                    if (protocol != null) {
                        res.addHeader(HttpHeaders.Names.WEBSOCKET_PROTOCOL, protocol);
                    }
                }

                ChannelPipeline p = ctx.getChannel().getPipeline();
                ctx.getChannel().write(res);
                // replace the HTTP decoder/encoder with WebSocketFrame
                // decoder/encoder
                p.replace("decoder", "wsdecoder", new WebSocketFrameDecoder());
                p.replace("encoder", "wsencoder", new WebSocketFrameEncoder());
                return;
            }
        }
        sendHttpResponse(ctx, req, new DefaultHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.FORBIDDEN));
    }

    /**
     * <p>
     * Responses the HTTP request (handshaking).
     * </p>
     * 
     * @param req
     *            the HttpRequest
     * @param ctx
     *            channel handler context, can't be null
     * @param res
     *            the HttpResponse
     */
    private void sendHttpResponse(ChannelHandlerContext ctx, HttpRequest req, HttpResponse res) {
        if (res.getStatus().getCode() != 200) {
            res.setContent(ChannelBuffers.copiedBuffer(res.getStatus().toString(), CharsetUtil.UTF_8));
            res.setHeader(HttpHeaders.Names.CONTENT_LENGTH, Integer.toString(res.getContent().readableBytes()));
        }

        ChannelFuture f = ctx.getChannel().write(res);
        if (!HttpHeaders.isKeepAlive(req) || res.getStatus().getCode() != 200) {
            f.addListener(ChannelFutureListener.CLOSE);
        }
    }

    /**
     * <p>
     * Method called when an exception has been thrown during channel operation.
     * Just prints the stack trace.
     * </p>
     * 
     * @param ctx
     *            channel handler context, can't be null
     * @param e
     *            message event, can't be null
     * @throws Exception
     *             if any error occurs
     */
    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, ExceptionEvent e) throws Exception {
        e.getCause().printStackTrace();
        e.getChannel().close();
    }

    /**
     * <p>
     * Return the web socket location.
     * </p>
     * 
     * @param req
     *            the HttpRequest
     * @return the web socket location.
     */
    private String getWebSocketLocation(HttpRequest req) {
        return "ws://" + req.getHeader(HttpHeaders.Names.HOST) + WEBSOCKET_PATH;
    }
}
